import datetime
import uuid

import torch
from diffusers import StableDiffusionPipeline
from fastapi import FastAPI, UploadFile, File, Form, Path, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import os, shutil
import json

from starlette.responses import HTMLResponse
from starlette.staticfiles import StaticFiles
from starlette.templating import Jinja2Templates
from pathlib import Path as SysPath
from train import router as train_router  # 引入 generate.py 中的 router
from generate import router as generate_router  # 引入 generate.py 中的 router
from search import router as search_router  # 引入 generate.py 中的 router
from upload import router as upload_router  # 引入 generate.py 中的 router

from inference import pipe


app = FastAPI()

# 目录初始化
UPLOAD_DIR = SysPath("./uploads")
UPLOAD_DIR.mkdir(exist_ok=True)
TEMPLATE_DIR = SysPath("./templates")
STATIC_DIR = SysPath("./static")
TEMPLATE_DIR.mkdir(exist_ok=True)
STATIC_DIR.mkdir(exist_ok=True)
GENERATED_DIR = SysPath("./generated")
GENERATED_DIR.mkdir(exist_ok=True)
MODEL_DIR = SysPath("./models")

# 注册静态资源目录
app.mount("/generated", StaticFiles(directory="./generated"), name="generated")
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.mount("/uploads", StaticFiles(directory=UPLOAD_DIR), name="uploads")

# 模板系统
templates = Jinja2Templates(directory=TEMPLATE_DIR)
# 概念存储
concepts_file = SysPath("./concepts.json")
if not concepts_file.exists():
    with open(concepts_file, "w") as f:
        json.dump({}, f)

# 首页
@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    uploads_dir = SysPath("./uploads")
    models_dir = SysPath("./models")

    trained_concepts = [p.name for p in models_dir.iterdir() if p.is_dir()]
    untrained_concepts = [
        p.name for p in uploads_dir.iterdir() if p.is_dir() and p.name not in trained_concepts
    ]

    return templates.TemplateResponse("index.html", {
        "request": request,
        "trained_concepts": trained_concepts,
        "untrained_concepts": untrained_concepts
    })
    # all_concepts = {d.name for d in UPLOAD_DIR.iterdir() if d.is_dir()}
    # trained_concepts = {f.stem for f in MODEL_DIR.glob("*.bin")}  # 你训练完会保存 .bin 或 .safetensors 等模型文件
    # untrained_concepts = all_concepts - trained_concepts
    #
    # return templates.TemplateResponse("index.html", {
    #     "request": request,
    #     "trained_concepts": sorted(trained_concepts),
    #     "untrained_concepts": sorted(untrained_concepts)
    # })
# 上传训练图像
app.include_router(upload_router)

# @app.post("/upload/")
# async def upload_image(concept: str = Form(...), file: UploadFile = File(...)):
#     concept_dir = UPLOAD_DIR / concept
#     concept_dir.mkdir(parents=True, exist_ok=True)
#     img_path = concept_dir / file.filename
#     with open(img_path, "wb") as f:
#         shutil.copyfileobj(file.file, f)
#
#     # 更新概念列表
#     with open(concepts_file) as f:
#         concepts = json.load(f)
#     if concept not in concepts:
#         concepts[concept] = []
#     concepts[concept].append(str(img_path))
#     with open(concepts_file, "w") as f:
#         json.dump(concepts, f)
#
#     return {"status": "ok", "path": str(img_path)}

# 训练接口
# @app.post("/train/")
# async def train(concept: str = Form(...)):
#     concept_dir = UPLOAD_DIR / concept
#     output_dir = SysPath("./models") / concept
#
#     if not concept_dir.exists():
#         return {"error": f"未找到概念: {concept}"}
#
#     success = train_concept(concept, concept_dir, output_dir)
#     return {"status": "ok" if success else "fail"}
# 注册训练接口路由
app.include_router(train_router)

# 注册图像生成路由
app.include_router(generate_router)

# 注册图像检索路由
app.include_router(search_router)